from __future__ import print_function, division
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
import torch
from SwinTransformer.models.swin_transformer import ViTSwinTransformerBlock, VitPatchMerging
import  BuildFormer.geoseg.models.BuildFormer as buildFormer
import UTNet.model.conv_trans_utils
class conv_block(nn.Module):
    """
    Convolution Block
    """

    def __init__(self, in_ch, out_ch):
        super(conv_block, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),

            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True))

    def forward(self, x):
        x = self.conv(x)

        return x


class up_conv(nn.Module):
    """
    Up Convolution Block
    """

    def __init__(self, in_ch, out_ch, scale_factor=2):
        super(up_conv, self).__init__()
        self.up = nn.Sequential(
            nn.Upsample(scale_factor=scale_factor),
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x
class conv_block1(nn.Module):
    """
    Convolution Block
    """

    def __init__(self, in_ch, out_ch):
        super().__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),

            nn.Conv2d(out_ch, out_ch, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True))

    def forward(self, x):
        x = self.conv(x)

        return x


class up_conv1(nn.Module):
    """
    Up Convolution Block
    """

    def __init__(self, in_ch, out_ch, scale_factor=2):
        super().__init__()
        self.up = nn.Sequential(
            nn.Upsample(scale_factor=scale_factor),
            nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x

class ConvBNRelu1D(nn.Module):
    def __init__(self, inChannel, outChannel, kernel_size=3, padding=1, bias=False):
        super(ConvBNRelu1D, self).__init__()
        self.conv = nn.Conv1d(inChannel, outChannel, kernel_size=kernel_size, padding=padding, bias=bias)
        self.bn = nn.BatchNorm1d(outChannel, momentum=0.1, affine=True)
        self.reLu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return self.reLu(x)


class LGFO(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch, hide_ch, out_ch, inWidth, hideWidth, outWidth):
        super().__init__()

        # self.layer1 =  nn.Linear(inWidth * inWidth, hideWidth * hideWidth)
        # self.CBR1=    ConvBNRelu1D(in_ch, hide_ch, kernel_size=1, padding=0, bias=True)
        #
        # self.layer2 = nn.Linear(hideWidth * hideWidth ,hideWidth * hideWidth)
        # self.CBR2 =    ConvBNRelu1D(hide_ch, hide_ch, kernel_size=1, padding=0, bias=True)
        #
        # self.layer3= nn.Linear(hideWidth * hideWidth, outWidth * outWidth)
        # self.CBR3 =    ConvBNRelu1D(hide_ch, out_ch, kernel_size=1, padding=0, bias=True)

        self.layer1= ConvBNRelu1D(in_ch, out_ch, kernel_size=1, padding=0, bias=True)
        self.layer2 =    nn.Linear(inWidth * inWidth, hideWidth * hideWidth)
         # ConvBNRelu1D(hide_ch, hide_ch, kernel_size=1, padding=0, bias=True),
        self.layer3 =    nn.Linear(hideWidth * hideWidth, hideWidth * hideWidth)
        self.layer4 =    nn.Linear(hideWidth * hideWidth, inWidth * inWidth)
        self.layer5 =    ConvBNRelu1D(out_ch, out_ch, kernel_size=1, padding=0, bias=True)


        self.outDim = out_ch

    def forward(self, x):
        # windows间交互
        b, c, w, h = x.shape
        x = x.view(b, c, w * h)

        x = self.layer1(x)
        shortcut = x
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = x + shortcut
        x = self.layer5(x)

        x = x.view(b, self.outDim, w, h)
        x = x.reshape(b, self.outDim, w, h)

        return x


class LGFO_Net(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        # self.Conv4=WinLearnedAttention(64,filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])
        # self.Conv5=WinLearnedAttention(32,filters[3], filters[4])

        self.Up5 = up_conv(filters[4], filters[3])
        self.Up_conv5 = conv_block(filters[4], filters[3])
        # self.Up_conv5 = WinLearnedAttention(64,filters[4], filters[3])

        self.Up4 = up_conv(filters[3], filters[2])
        self.Up_conv4 = conv_block(filters[3], filters[2])

        self.Up3 = up_conv(filters[2], filters[1])
        self.Up_conv3 = conv_block(filters[2], filters[1])

        self.Up2 = up_conv(filters[1], filters[0])
        self.Up_conv2 = conv_block(filters[1], filters[0])

        self.Conv = nn.Conv2d(filters[0], out_ch, kernel_size=1, stride=1, padding=0)

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        d5 = self.Up5(e5)
        d5 = torch.cat((e4, d5), dim=1)

        d5 = self.Up_conv5(d5)

        d4 = self.Up4(d5)
        d4 = torch.cat((e3, d4), dim=1)
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        d3 = torch.cat((e2, d3), dim=1)
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        d2 = torch.cat((e1, d2), dim=1)
        d2 = self.Up_conv2(d2)

        out = self.Conv(d2)

        # d1 = self.active(out)

        return out
class LGFO_Net_NoConv(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)


        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)


        self.UNetDecoder = UnetDecoder(depths=filters)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        out=self.UNetDecoder([g1,g2,g3,g4,g5])

        return out

class mutiDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose1_0 = up_conv(depths[1], depths[0])
        self.Conv1_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose2_1 = up_conv(depths[2], depths[1])
        self.Conv2_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose2_0 = up_conv(depths[1], depths[0])
        self.Conv2_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose3_2 = up_conv(depths[3], depths[2])
        self.Conv3_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose3_1 = up_conv(depths[2], depths[1])
        self.Conv3_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose3_0 = up_conv(depths[1], depths[0])
        self.Conv3_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3])
        self.Conv4_3 = conv_block(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv(depths[3], depths[2])
        self.Conv4_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv(depths[2], depths[1])
        self.Conv4_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv(depths[1], depths[0])
        self.Conv4_0 = conv_block(depths[0], depths[0])

        self.ConvAll = conv_block(depths[0] * 5, depths[0])
        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        level0 = levels[0]
        # level1的解码
        level1 = self.ConvTranspose1_0(levels[1])
        level1 = self.Conv1_0(level1)

        # level2的解码
        level2 = self.ConvTranspose2_1(levels[2])
        level2 = torch.cat([level2, levels[1]], dim=1)
        level2 = self.Conv2_1(level2)

        level2 = self.ConvTranspose2_0(level2)
        level2 = self.Conv2_0(level2)

        # level3的解码
        level3 = self.ConvTranspose3_2(levels[3])
        level3 = torch.cat([level3, levels[2]], dim=1)
        level3 = self.Conv3_2(level3)

        level3 = self.ConvTranspose3_1(level3)
        level3 = torch.cat([level3, levels[1]], dim=1)
        level3 = self.Conv3_1(level3)

        level3 = self.ConvTranspose3_0(level3)
        level3 = self.Conv3_0(level3)

        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        # 所有level的融合
        out = torch.cat([level0, level1, level2, level3, level4], dim=1)
        out = self.ConvAll(out)
        out = self.Out(out)

        return out


class UnetDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3])
        self.Conv4_3 = conv_block(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv(depths[3], depths[2])
        self.Conv4_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv(depths[2], depths[1])
        self.Conv4_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv(depths[1], depths[0])
        self.Conv4_0 = conv_block(depths[0], depths[0])

        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        out = self.Out(level4)

        return out
class UnetDecoder1(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose4_3 = up_conv1(depths[4], depths[3])
        self.Conv4_3 = conv_block1(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv1(depths[3], depths[2])
        self.Conv4_2 = conv_block1(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv1(depths[2], depths[1])
        self.Conv4_1 = conv_block1(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv1(depths[1], depths[0])
        self.Conv4_0 = conv_block1(depths[0], depths[0])

        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        out = self.Out(level4)

        return out

class DenseDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose1_0 = up_conv(depths[1], depths[0], scale_factor=2)

        self.ConvTranspose2_1 = up_conv(depths[2], depths[1], scale_factor=2)

        self.ConvTranspose2_0 = up_conv(depths[2], depths[0], scale_factor=4)

        self.ConvTranspose3_2 = up_conv(depths[3], depths[2], scale_factor=2)

        self.ConvTranspose3_1 = up_conv(depths[3], depths[1], scale_factor=4)

        self.ConvTranspose3_0 = up_conv(depths[3], depths[0], scale_factor=8)

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3], scale_factor=2)

        self.ConvTranspose4_2 = up_conv(depths[4], depths[2], scale_factor=4)

        self.ConvTranspose4_1 = up_conv(depths[4], depths[1], scale_factor=8)

        self.ConvTranspose4_0 = up_conv(depths[4], depths[0], scale_factor=16)

        self.upConv3 = conv_block(depths[3] * 2, depths[3])
        self.upConv2 = conv_block(depths[2] * 3, depths[2])
        self.upConv1 = conv_block(depths[1] * 4, depths[1])
        self.upConv0 = conv_block(depths[0] * 5, depths[0])

        self.out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        l4_3 = self.ConvTranspose4_3(levels[4])
        l4_2 = self.ConvTranspose4_2(levels[4])
        l4_1 = self.ConvTranspose4_1(levels[4])
        l4_0 = self.ConvTranspose4_0(levels[4])

        levels[3] = torch.cat([levels[3], l4_3], dim=1)
        levels[3] = self.upConv3(levels[3])

        l3_2 = self.ConvTranspose3_2(levels[3])
        l3_1 = self.ConvTranspose3_1(levels[3])
        l3_0 = self.ConvTranspose3_0(levels[3])

        levels[2] = torch.cat([levels[2], l4_2, l3_2], dim=1)
        levels[2] = self.upConv2(levels[2])

        l2_1 = self.ConvTranspose2_1(levels[2])
        l2_0 = self.ConvTranspose2_0(levels[2])

        levels[1] = torch.cat([levels[1], l4_1, l3_1, l2_1], dim=1)
        levels[1] = self.upConv1(levels[1])

        l1_0 = self.ConvTranspose1_0(levels[1])

        levels[0] = torch.cat([levels[0], l4_0, l3_0, l2_0, l1_0], dim=1)
        levels[0] = self.upConv0(levels[0])

        out = self.out(levels[0])

        return out


class LGFO_Net1(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        # self.Conv4=WinLearnedAttention(64,filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])
        # self.Conv5=WinLearnedAttention(32,filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

        self.decoder = decoder(filters)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.decoder([e1, e2, e3, e4, e5])

        # d1 = self.active(out)

        return out


class LGFO_Net_denseDecode(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        # self.Conv4=WinLearnedAttention(64,filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])
        # self.Conv5=WinLearnedAttention(32,filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

        self.decoder = DenseDecoder(filters)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.decoder([e1, e2, e3, e4, e5])

        # d1 = self.active(out)

        return out


# use swintransformerblock and LGFO architecture
class LGFO_SwinTransformerNet(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = nn.Sequential(
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=4)
            )
        self.interWin2 = nn.Sequential(VitPatchMerging(filters[0]),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin3 = nn.Sequential(VitPatchMerging(filters[1]),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin4 = nn.Sequential(VitPatchMerging(filters[2]),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin5 = nn.Sequential(VitPatchMerging(filters[3]),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=4)
                                       )

        self.UNetDecoder = UnetDecoder(filters)

    def forward(self, input):
        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        g1 = self.interWin1(e1)
        g2 = self.interWin2(g1)
        g3 = self.interWin3(g2)
        g4 = self.interWin4(g3)
        g5 = self.interWin5(g4)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.UNetDecoder([e1, e2, e3, e4, e5])

        return out


class SwinTransformer_UNetDecoder(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Conv1 = conv_block(in_ch, filters[0])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = nn.Sequential(
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=4)
            )
        self.interWin2 = nn.Sequential(VitPatchMerging(filters[0]),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin3 = nn.Sequential(VitPatchMerging(filters[1]),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin4 = nn.Sequential(VitPatchMerging(filters[2]),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin5 = nn.Sequential(VitPatchMerging(filters[3]),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=4)
                                       )

        self.UNetDecoder = UnetDecoder(filters)

    def forward(self, input):
        e1 = self.Conv1(input)

        g1 = self.interWin1(e1)
        g2 = self.interWin2(g1)
        g3 = self.interWin3(g2)
        g4 = self.interWin4(g3)
        g5 = self.interWin5(g4)

        out = self.UNetDecoder([g1, g2, g3, g4, g5])

        return out


class layer(nn.Module):
    def __init__(self, size=(256,256), inDim=32,num_heads=1):
        super().__init__()


        self.size=size
        self.layer1=ViTSwinTransformerBlock(size, dim=inDim, num_heads=num_heads, window_size=8, shift_size=0)
        self.layer2=ViTSwinTransformerBlock([8, 8], dim=inDim, num_heads=num_heads, window_size=8, shift_size=0)
        self.layer3 = ViTSwinTransformerBlock(size, dim=inDim, num_heads=num_heads, window_size=8, shift_size=0)
        self.layer4 =ViTSwinTransformerBlock([8, 8], dim=inDim, num_heads=num_heads, window_size=8, shift_size=0)
    def forward(self, x):

        x=self.layer1(x)
        short1 =x
        x=F.interpolate(x,size=(8,8))
        x=self.layer2(x)
        x=F.interpolate(x,size=self.size)+short1

        x = self.layer3(x)
        short3 = x
        x = F.interpolate(x, size=(8, 8))
        x = self.layer4(x)
        x = F.interpolate(x, size=self.size) + short3

        return x

# use not shift win transformerblock and LGFO architecture
class NSwinTransformer(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = nn.Sequential(
            layer([256, 256], filters[0],num_heads=1),
            )
        self.interWin2 = nn.Sequential(VitPatchMerging(filters[0]),
                                       layer([128, 128], filters[1],num_heads=2),
                                       )
        self.interWin3 = nn.Sequential(VitPatchMerging(filters[1]),
                                       layer([64, 64], filters[2],num_heads=4),
                                       )
        self.interWin4 = nn.Sequential(VitPatchMerging(filters[2]),
                                       layer([32, 32], filters[3],num_heads=8),
                                       )
        self.interWin5 = nn.Sequential(VitPatchMerging(filters[3]),
                                       layer([16, 16], filters[4],num_heads=16),
                                       )

        self.UNetDecoder = UnetDecoder(filters)

    def forward(self, input):
        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        g1 = self.interWin1(e1)
        g2 = self.interWin2(g1)
        g3 = self.interWin3(g2)
        g4 = self.interWin4(g3)
        g5 = self.interWin5(g4)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.UNetDecoder([e1, e2, e3, e4, e5])

        return out

class testLayer(nn.Module):
    def __init__(self, inDim=32, outDim=32,num_heads=1):
        super().__init__()

        # self.preprocess = nn.Conv2d(inDim, outDim, kernel_size=1, stride=1, padding=0)
        self.layer1=conv_block(inDim, outDim)
        # self.CNN1 = nn.Conv2d(inDim, inDim, kernel_size=3, stride=1, padding=1)
        # self.CNN1 = buildFormer.LWMSA(dim=inDim, num_heads=num_heads, window_size=8)
        # self.layer2=ViTSwinTransformerBlock([8, 8], dim=outDim, num_heads=num_heads, window_size=8, shift_size=0)
        self.layer2=nn.Sequential(
            conv_block(outDim, outDim),

        )
        # nn.Conv2d(outDim, outDim, kernel_size=3, stride=1, padding=1)
        # self.layer3 = conv_block(outDim, outDim)
        # self.CNN2 = buildFormer.LWMSA(dim=inDim, num_heads=num_heads, window_size=8)
        self.relu=nn.ReLU()
    def forward(self, x):
        _,_,h,w=x.size()
        # x=self.preprocess(x)
        x=self.layer1(x)
        short1 =x
        x=F.interpolate(x,size=(8,8))
        x=self.layer2(x)
        x=F.interpolate(x,size=(h,w))+short1
        # x=self.relu(x)



        return x
def merge_channels(input_channels, output_channels, image):
    # 检查输入通道数能否被输出通道数整除
    if input_channels % output_channels != 0:
        print("Error: 输入通道数不能被输出通道数整除")
        return None

    # 计算每个输出通道需要合并的输入通道数
    channels_per_group = input_channels // output_channels

    # 将输入图像沿最后一个维度分割成输入通道数个通道
    image_channels = torch.split(image, 1, dim=-1)

    # 创建空列表存储合并后的通道
    merged_channels = []

    # 遍历输出通道数
    for i in range(output_channels):
        # 创建一个全零张量
        merged_channel = torch.zeros_like(image_channels[0])
        # 对于每个输出通道，将相邻的输入通道相加
        for j in range(channels_per_group):
            index = i * channels_per_group + j
            merged_channel = merged_channel+image_channels[index]
        # 将合并后的通道添加到列表中
        merged_channels.append(merged_channel)

    # 将合并后的通道沿最后一个维度连接
    merged_image = torch.cat(merged_channels, dim=-1)
    return merged_image


class DepthShareCBN(nn.Module):
    """
    Convolution Block
    """

    def __init__(self, in_ch, out_ch):
        super().__init__()

        self.Conv =nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=True)
        self.PointConv =nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1, padding=0, bias=True)
        self.BN = nn.BatchNorm2d(out_ch)
        self.Relu= nn.ReLU(inplace=True)



    def forward(self, x):
        B,C,H,W=x.size()
        x=x.reshape(-1,1,H,W)
        x = self.Conv(x)
        x=x.reshape(B,C,H,W)
        x=self.PointConv(x)
        x=self.BN(x)
        x=self.Relu(x)
        return x
class test(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.local = nn.Sequential(
            conv_block(in_ch, filters[0]),
            conv_block(filters[0], filters[0]),
            conv_block(filters[0], filters[0]),
            conv_block(filters[0], filters[0]),
            conv_block(filters[0], filters[0]),
            # conv_block(filters[0], filters[0]),
            # conv_block(filters[0], filters[0]),
            # conv_block(filters[0], filters[0]),
            # conv_block(filters[0], filters[0]),
            # conv_block(filters[0], filters[0])
        )

        # self.Conv1 = conv_block(in_ch, 64)


        # self.Conv2 = nn.Sequential(
        #     conv_block(filters[0], filters[1]),
        #     conv_block(filters[1], filters[1])
        #                            )
        # self.Conv3 = nn.Sequential(
        #     conv_block(filters[1], filters[2]),
        #     conv_block(filters[2], filters[2])
        #                            )
        # self.Conv4 = nn.Sequential(
        #     conv_block(filters[2], filters[3]),
        #     conv_block(filters[3], filters[3])
        #                            )
        # self.Conv5 = nn.Sequential(
        #     conv_block(filters[3], filters[4]),
        #     conv_block(filters[4], filters[4])
        #                            )




        # self.Trans1=nn.Sequential(
        #     ViTSwinTransformerBlock([256,256],dim=64,num_heads=1,window_size=8,shift_size=0),
        #     ViTSwinTransformerBlock([256, 256], dim=64, num_heads=1, window_size=8, shift_size=4),
        #     )
        # self.Trans2 = nn.Sequential(
        #     ViTSwinTransformerBlock([128, 128], dim=64, num_heads=2, window_size=8, shift_size=0),
        #     ViTSwinTransformerBlock([128, 128], dim=64, num_heads=2, window_size=8, shift_size=4),
        # )
        # self.Trans3 = nn.Sequential(
        #     ViTSwinTransformerBlock([64, 64], dim=64, num_heads=4, window_size=8, shift_size=0),
        #     ViTSwinTransformerBlock([64, 64], dim=64, num_heads=4, window_size=8, shift_size=4),
        # )
        self.Trans4 = nn.Sequential(
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
        )
        self.Trans5 = nn.Sequential(
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),

        )

        # self.Trans = nn.Sequential(
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=0),
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=4),
        #
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=0),
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=4),
        # )


        # self.mutiDecoder = mutiDecoder(filters)
        # self.UnetDecoder = UnetDecoder(filters)
        # self.Conv64_2 = nn.Conv2d(64, 2, kernel_size=1, stride=1, padding=0)
        #
        # self.Conv128_64=nn.Conv2d(128,64,kernel_size=1,stride=1,padding=0)
        # self.Conv256_64=nn.Conv2d(256,64,kernel_size=1,stride=1,padding=0)
        # self.Conv512_64=nn.Conv2d(512,64,kernel_size=1,stride=1,padding=0)
        # self.Conv1024_64=nn.Conv2d(1024,64,kernel_size=1,stride=1,padding=0)
        # #
        # self.ConvOut1 = nn.Conv2d(64, 16 , kernel_size=1, stride=1, padding=0)
        # self.ConvOut2 = nn.Conv2d(16 , 2, kernel_size=1, stride=1, padding=0)
        self.UnetDecoder=UnetDecoder([64,64,64,64,64])
        # self.MutiDecoder=mutiDecoder([64,64,64,64,64])
    def forward(self, input):

        e1=self.local(input)

        e2=F.interpolate(input,size=(128,128))
        e2= self.local(e2)

        e3=F.interpolate(input,size=(64,64))
        e3= self.local(e3)

        # e4=F.interpolate(input,size=(32,32))
        # e4= self.local(e4)
        # #
        # e5=F.interpolate(input,size=(16,16))
        # e5= self.local(e5)


        # local = self.Conv(input)
        #

        e4=self.Maxpool1(e3)
        e4=self.Trans4(e4)


        e5=self.Maxpool4(e4)
        e5=self.Trans5(e5)


        #
        # g1 = self.Conv64_2(g1)
        #
        # out=F.interpolate(g1,size=(256,256))
        #
        #
        # # out=torch.cat([local,g1],dim=1)
        #
        #
        # # out=self.ConvOut1(out)
        # # out=self.ConvOut2(out)


        out = self.UnetDecoder([e1, e2, e3, e4, e5])

        return out
class test1(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        # self.local = nn.Sequential(
        #     conv_block(in_ch, filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0])
        # )
        self.conv=conv_block(3, filters[0])
        self.local = nn.Sequential(
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=8, shift_size=4),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=8, shift_size=4),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=8, shift_size=4),
        )
        # self.Conv1 = conv_block(in_ch, 64)


        # self.Conv2 = nn.Sequential(
        #     conv_block(filters[0], filters[1]),
        #     conv_block(filters[1], filters[1])
        #                            )
        # self.Conv3 = nn.Sequential(
        #     conv_block(filters[1], filters[2]),
        #     conv_block(filters[2], filters[2])
        #                            )
        # self.Conv4 = nn.Sequential(
        #     conv_block(filters[2], filters[3]),
        #     conv_block(filters[3], filters[3])
        #                            )
        # self.Conv5 = nn.Sequential(
        #     conv_block(filters[3], filters[4]),
        #     conv_block(filters[4], filters[4])
        #                            )




        # self.Trans1=nn.Sequential(
        #     ViTSwinTransformerBlock([256,256],dim=64,num_heads=1,window_size=8,shift_size=0),
        #     ViTSwinTransformerBlock([256, 256], dim=64, num_heads=1, window_size=8, shift_size=4),
        #     )
        # self.Trans2 = nn.Sequential(
        #     ViTSwinTransformerBlock([128, 128], dim=64, num_heads=2, window_size=8, shift_size=0),
        #     ViTSwinTransformerBlock([128, 128], dim=64, num_heads=2, window_size=8, shift_size=4),
        # )
        # self.Trans3 = nn.Sequential(
        #     ViTSwinTransformerBlock([64, 64], dim=64, num_heads=4, window_size=8, shift_size=0),
        #     ViTSwinTransformerBlock([64, 64], dim=64, num_heads=4, window_size=8, shift_size=4),
        # )
        self.Trans4 = nn.Sequential(
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
        )
        self.Trans5 = nn.Sequential(
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),

        )

        # self.Trans = nn.Sequential(
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=0),
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=4),
        #
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=0),
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=4),
        # )


        # self.mutiDecoder = mutiDecoder(filters)
        # self.UnetDecoder = UnetDecoder(filters)
        # self.Conv64_2 = nn.Conv2d(64, 2, kernel_size=1, stride=1, padding=0)
        #
        # self.Conv128_64=nn.Conv2d(128,64,kernel_size=1,stride=1,padding=0)
        # self.Conv256_64=nn.Conv2d(256,64,kernel_size=1,stride=1,padding=0)
        # self.Conv512_64=nn.Conv2d(512,64,kernel_size=1,stride=1,padding=0)
        # self.Conv1024_64=nn.Conv2d(1024,64,kernel_size=1,stride=1,padding=0)
        # #
        # self.ConvOut= nn.Conv2d(64*5, 64 , kernel_size=1, stride=1, padding=0)
        # self.ConvOut2 = nn.Conv2d(64 , 2, kernel_size=1, stride=1, padding=0)
        self.UnetDecoder=UnetDecoder([64,64,64,64,64])
        # self.MutiDecoder=mutiDecoder([64,64,64,64,64])
    def forward(self, input):

        input=self.conv(input)
        e1=self.local(input)

        e2=F.interpolate(input,size=(128,128))
        e2= self.local(e2)

        e3=F.interpolate(input,size=(64,64))
        e3= self.local(e3)

        # e4=F.interpolate(input,size=(32,32))
        # e4= self.local(e4)
        # #
        # e5=F.interpolate(input,size=(16,16))
        # e5= self.local(e5)


        # local = self.Conv(input)
        #

        e4=self.Maxpool1(e3)
        e4=self.Trans4(e4)


        e5=self.Maxpool4(e4)
        e5=self.Trans5(e5)


        #
        # g1 = self.Conv64_2(g1)
        #
        # out=F.interpolate(g1,size=(256,256))
        #
        #
        # # out=torch.cat([local,g1],dim=1)
        #
        #
        # # out=self.ConvOut1(out)
        # # out=self.ConvOut2(out)

        # e2=F.interpolate(e2,size=(256,256))
        # e3=F.interpolate(e3,size=(256,256))
        # e4=F.interpolate(e4,size=(256,256))
        # e5=F.interpolate(e5,size=(256,256))
        # out=torch.cat([e1,e2,e3,e4,e5],dim=1)
        # out=self.ConvOut(out)
        # out = self.ConvOut2(out)
        out = self.UnetDecoder([e1, e2, e3, e4, e5])

        return out
class test2(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        # self.local = nn.Sequential(
        #     conv_block(in_ch, filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0]),
        #     # conv_block(filters[0], filters[0])
        # )
        self.conv1=conv_block(3, filters[0])
        self.conv2 = conv_block(filters[0], filters[0])
        self.conv3 = conv_block(filters[0], filters[0])

        # self.Conv1 = conv_block(in_ch, 64)


        # self.Conv2 = nn.Sequential(
        #     conv_block(filters[0], filters[1]),
        #     conv_block(filters[1], filters[1])
        #                            )
        # self.Conv3 = nn.Sequential(
        #     conv_block(filters[1], filters[2]),
        #     conv_block(filters[2], filters[2])
        #                            )
        # self.Conv4 = nn.Sequential(
        #     conv_block(filters[2], filters[3]),
        #     conv_block(filters[3], filters[3])
        #                            )
        # self.Conv5 = nn.Sequential(
        #     conv_block(filters[3], filters[4]),
        #     conv_block(filters[4], filters[4])
        #                            )




        # self.Trans1=nn.Sequential(
        #     ViTSwinTransformerBlock([256,256],dim=64,num_heads=1,window_size=8,shift_size=0),
        #     ViTSwinTransformerBlock([256, 256], dim=64, num_heads=1, window_size=8, shift_size=4),
        #     )
        # self.Trans2 = nn.Sequential(
        #     ViTSwinTransformerBlock([128, 128], dim=64, num_heads=2, window_size=8, shift_size=0),
        #     ViTSwinTransformerBlock([128, 128], dim=64, num_heads=2, window_size=8, shift_size=4),
        # )
        # self.Trans3 = nn.Sequential(
        #     ViTSwinTransformerBlock([64, 64], dim=64, num_heads=4, window_size=8, shift_size=0),
        #     ViTSwinTransformerBlock([64, 64], dim=64, num_heads=4, window_size=8, shift_size=4),
        # )
        self.Trans4 = nn.Sequential(
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
            ViTSwinTransformerBlock([32, 32], dim=64, num_heads=8, window_size=32, shift_size=0),
        )
        self.Trans5 = nn.Sequential(
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),
            ViTSwinTransformerBlock([16, 16], dim=64, num_heads=8, window_size=16, shift_size=0),

        )

        # self.Trans = nn.Sequential(
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=0),
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=4),
        #
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=0),
        #     ViTSwinTransformerBlock([16, 16], dim=64, num_heads=16, window_size=16, shift_size=4),
        # )


        # self.mutiDecoder = mutiDecoder(filters)
        # self.UnetDecoder = UnetDecoder(filters)
        # self.Conv64_2 = nn.Conv2d(64, 2, kernel_size=1, stride=1, padding=0)
        #
        # self.Conv128_64=nn.Conv2d(128,64,kernel_size=1,stride=1,padding=0)
        # self.Conv256_64=nn.Conv2d(256,64,kernel_size=1,stride=1,padding=0)
        # self.Conv512_64=nn.Conv2d(512,64,kernel_size=1,stride=1,padding=0)
        # self.Conv1024_64=nn.Conv2d(1024,64,kernel_size=1,stride=1,padding=0)
        # #
        # self.ConvOut1 = nn.Conv2d(64, 16 , kernel_size=1, stride=1, padding=0)
        # self.ConvOut2 = nn.Conv2d(16 , 2, kernel_size=1, stride=1, padding=0)
        self.UnetDecoder=UnetDecoder([64,64,64,64,64])
        # self.MutiDecoder=mutiDecoder([64,64,64,64,64])
    def forward(self, input):

        e1=self.conv1(input)

        e2=self.Maxpool1(e1)
        e2= self.conv2(e2)

        e3=self.Maxpool2(e2)
        e3= self.conv3(e3)

        # e4=F.interpolate(input,size=(32,32))
        # e4= self.local(e4)
        # #
        # e5=F.interpolate(input,size=(16,16))
        # e5= self.local(e5)


        # local = self.Conv(input)
        #

        e4=self.Maxpool3(e3)
        e4=self.Trans4(e4)


        e5=self.Maxpool4(e4)
        e5=self.Trans5(e5)


        #
        # g1 = self.Conv64_2(g1)
        #
        # out=F.interpolate(g1,size=(256,256))
        #
        #
        # # out=torch.cat([local,g1],dim=1)
        #
        #
        # # out=self.ConvOut1(out)
        # # out=self.ConvOut2(out)


        out = self.UnetDecoder([e1, e2, e3, e4, e5])

        return out